from .model import Model
from .refer_path import ReferPath
import cvxpy
import numpy as np

# from model import Model
# from refer_path import ReferPath


class MPC:
    def __init__(
        self,
        model: Model,
        path: ReferPath,
        T: int,
        timeStep: float,
        Nx: int,
        Nu: int,
        Q: list[float],
        R: list[float],
    ) -> None:
        self.model = model
        self.path = path
        self.T = T
        self.timeStep = timeStep
        self.nx = Nx
        self.nu = Nu
        self.Q = np.diag(Q)
        self.R = np.diag(R)

    def is_ready(self) -> bool:
        return self.path.is_ready()

    def getControl(
        self, current: list[float]
    ) -> tuple[list[float], list[list[float]], list[list[float]]] | None:
        """
        current: list[float]  [theta,v,x,y]
        return: tuple[list[float], list[list[float]]]
            control: list[float]  [w,a]
            states: list[list[float]]  [[theta,v,x,y],...]
            controls: list[list[float]]  [[w,a],...]
        """
        if not self.path.is_ready():
            return [0, 0], [], []
        # creat referance path
        referStates, referControl = self.path.getReferancePath(
            current, self.T, self.timeStep
        )

        # solve mpc problem
        x = cvxpy.Variable((self.nx, self.T + 1))
        u = cvxpy.Variable((self.nu, self.T))

        cost = 0.0
        constraints = []

        for t in range(self.T):
            cost += cvxpy.quad_form(u[:, t] - referControl[:, t], self.R)
            if t != 0:
                cost += cvxpy.quad_form(x[:, t] - referStates[:, t], self.Q)
            matA = self.model.getMatrixA(
                referStates[1, t], referStates[0, t], self.timeStep
            )
            matB = self.model.getMatrixB(self.timeStep)
            constraints += [
                x[:, t + 1] - referStates[:, t + 1]
                == matA @ (x[:, t] - referStates[:, t])
                + matB @ (u[:, t] - referControl[:, t])
            ]
        cost += cvxpy.quad_form(x[:, self.T] - referStates[:, self.T], self.Q)

        constraints += [(x[0, 0]) == current[0]]
        constraints += [(x[1, 0]) == current[1]]
        constraints += [(x[2, 0]) == current[2]]
        constraints += [(x[3, 0]) == current[3]]
        constraints += [cvxpy.abs(u[0, :]) <= self.model.wRange[1]]
        constraints += [cvxpy.abs(u[1, :]) <= self.model.aRange[1]]
        constraints += [cvxpy.abs(x[1, :]) <= self.model.vRange[1]]

        prob = cvxpy.Problem(cvxpy.Minimize(cost), constraints)
        prob.solve(verbose=False)

        if prob.status == cvxpy.OPTIMAL or prob.status == cvxpy.OPTIMAL_INACCURATE:
            if x.value is not None and u.value is not None:
                a = u.value[1, 0]
                w = u.value[0, 0]
                states = []
                for i in range(self.T + 1):
                    states.append(
                        [x.value[0, i], x.value[1, i], x.value[2, i], x.value[3, i]]
                    )
                controls = []
                for i in range(self.T):
                    controls.append([u.value[0, i], u.value[1, i]])
                return [w, a], states, controls
        print("Error: Cannot solve mpc..")
        return None

    def drawPrediction(
        self, ax, states, color: str, index: bool = False, points: bool = False
    ):
        import matplotlib.pyplot as plt

        x = [s[2] for s in states]
        y = [s[3] for s in states]
        ax.plot(x, y, color=color)
        if points:
            ax.scatter(x, y, color=color)
        if index:
            for i in range(len(states)):
                ax.text(x[i], y[i], str(i))


if __name__ == "__main__":
    from refer_path import Spline
    from refer_path import ReferPath
    import matplotlib.pyplot as plt

    # line = Spline([(-1, 0), (0, 0), (1, 0), (2, 1), (3, 0), (4, -1), (5, 0)], 1.0, 0.02)
    line = Spline([(0, 0), (-10, 0), (10, 10), (-20, 10)], 1.0, 0.02)
    T = 10
    timeStep = 0.1
    path = ReferPath(line)
    model = Model((-100, 100), (-100, 100), (-100, 100))
    mpc = MPC(
        model, path, T=T, timeStep=timeStep, Nx=4, Nu=2, Q=[10, 1, 10, 10], R=[1, 1]
    )
    curCtrl, state, control = mpc.getControl([0, 0, 0, 0])

    ax = plt.subplot(1, 2, 1)
    line.drawPath(ax, "blue")
    mpc.drawPrediction(ax, state, "red", True, True)
    referState, referControl = path.getReferancePath((0, 0), T, timeStep)
    path.drawReferancePath(referState, ax, "green")

    bx = plt.subplot(1, 2, 2)
    bx.plot([s[0] for s in state], label="theta")
    bx.plot([s[1] for s in state], label="v")
    bx.plot([s[2] for s in state], label="x")
    bx.plot([s[3] for s in state], label="y")
    bx.plot([c[0] for c in control], label="w")
    bx.plot([c[1] for c in control], label="a")
    bx.legend()
    plt.show()
